package top.conangao.common.security.handler;

import lombok.extern.slf4j.Slf4j;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpStatus;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.authority.SimpleGrantedAuthority;
import org.springframework.security.oauth2.core.OAuth2AuthenticatedPrincipal;
import org.springframework.security.oauth2.server.resource.introspection.BadOpaqueTokenException;
import org.springframework.security.oauth2.server.resource.introspection.OAuth2IntrospectionAuthenticatedPrincipal;
import org.springframework.security.oauth2.server.resource.introspection.OAuth2IntrospectionException;
import org.springframework.security.oauth2.server.resource.introspection.ReactiveOpaqueTokenIntrospector;
import org.springframework.util.Assert;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.client.ClientResponse;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Mono;

import java.net.URI;
import java.time.Instant;
import java.util.*;

/**
 * @author ConanGao
 * @description
 * @since 2024/2/22 10:33
 **/
@Slf4j
public class CustomReactiveOpaqueTokenIntrospector implements ReactiveOpaqueTokenIntrospector {
    private static final String AUTHORITY_PREFIX = "SCOPE_";
    private static final ParameterizedTypeReference<Map<String, Object>> STRING_OBJECT_MAP = new ParameterizedTypeReference<>() {
    };
    private final URI introspectionUri;
    private final WebClient webClient;

    public CustomReactiveOpaqueTokenIntrospector(String introspectionUri, String clientId, String clientSecret) {
        Assert.hasText(introspectionUri, "introspectionUri cannot be empty");
        Assert.hasText(clientId, "clientId cannot be empty");
        Assert.notNull(clientSecret, "clientSecret cannot be null");
        this.introspectionUri = URI.create(introspectionUri);
        this.webClient = WebClient.builder().defaultHeaders(h -> h.setBasicAuth(clientId, clientSecret)).build();
    }

    public CustomReactiveOpaqueTokenIntrospector(String introspectionUri, WebClient webClient) {
        Assert.hasText(introspectionUri, "introspectionUri cannot be null");
        Assert.notNull(webClient, "webClient cannot be null");
        this.introspectionUri = URI.create(introspectionUri);
        this.webClient = webClient;
    }

    public Mono<OAuth2AuthenticatedPrincipal> introspect(String token) {
        return Mono.just(token).flatMap(this::makeRequest).flatMap(this::adaptToNimbusResponse).map(this::convertClaimsSet).onErrorMap((e) -> !(e instanceof OAuth2IntrospectionException), this::onError);
    }

    private Mono<ClientResponse> makeRequest(String token) {
        return this.webClient.post().uri(this.introspectionUri).header("Accept", new String[]{"application/json"}).body(BodyInserters.fromFormData("token", token)).exchange();
    }

    private Mono<Map<String, Object>> adaptToNimbusResponse(ClientResponse responseEntity) {
        return responseEntity.statusCode() != HttpStatus.OK ? responseEntity.bodyToFlux(DataBuffer.class).map(DataBufferUtils::release).then(Mono.error(new OAuth2IntrospectionException("Introspection endpoint responded with " + responseEntity.statusCode()))) : responseEntity.bodyToMono(STRING_OBJECT_MAP).filter((body) -> (Boolean) body.compute("active", (k, v) -> {
            if (v instanceof String) {
                return Boolean.parseBoolean((String) v);
            } else {
                return v instanceof Boolean ? v : false;
            }
        })).switchIfEmpty(Mono.error(() -> new BadOpaqueTokenException("Provided token isn't active")));
    }

    private OAuth2AuthenticatedPrincipal convertClaimsSet(Map<String, Object> claims) {
        claims.computeIfPresent("aud", (k, v) -> v instanceof String ? Collections.singletonList(v) : v);
        claims.computeIfPresent("client_id", (k, v) -> v.toString());
        claims.computeIfPresent("exp", (k, v) -> Instant.ofEpochSecond(((Number) v).longValue()));
        claims.computeIfPresent("iat", (k, v) -> Instant.ofEpochSecond(((Number) v).longValue()));
        claims.computeIfPresent("iss", (k, v) -> v.toString());
        claims.computeIfPresent("nbf", (k, v) -> Instant.ofEpochSecond(((Number) v).longValue()));
        Collection<GrantedAuthority> authorities = new ArrayList<>();
        claims.computeIfPresent("scope", (k, v) -> {
            if (!(v instanceof String)) {
                return v;
            } else {
                Collection<String> scopes = Arrays.asList(((String) v).split(" "));

                for (String scope : scopes) {
                    authorities.add(new SimpleGrantedAuthority("SCOPE_" + scope));
                }

                return scopes;
            }
        });
        return new OAuth2IntrospectionAuthenticatedPrincipal(claims, authorities);
    }

    private OAuth2IntrospectionException onError(Throwable ex) {
        return new OAuth2IntrospectionException(ex.getMessage(), ex);
    }
}
